import pandas
from pandas import datetime
from statsmodels.tsa.stattools import adfuller
from statsmodels.stats.diagnostic import acorr_ljungbox #白噪声检验
from statsmodels.graphics.tsaplots import plot_acf
from statsmodels.graphics.tsaplots import plot_pacf
from matplotlib import pyplot
import numpy as np
from pandas import read_csv
from sklearn.metrics import mean_squared_error
from statsmodels.tsa.arima_model import ARIMA
from math import sqrt
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

def parse(x):
    return datetime.strptime(x, '%Y %m')

def mean_absolute_percentage_error(y_true, y_pred):
    # 平均绝对百分比误差（MAPE）的计算
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    return np.mean(np.abs((y_true - y_pred) / y_true)) * 100


# evaluate the RMSE for each forecast time step
def evaluate_forecasts(test, forecasts):
    rmse = sqrt(mean_squared_error(test, forecasts))
    mae = mean_absolute_error(test, forecasts)
    mape = mean_absolute_percentage_error(test, forecasts)
    r2 = r2_score(test, forecasts)
    print(' RMSE: %f' % rmse)
    print(' MAE: %f' % mae)
    print(' MAPE: %f' % mape)
    print(' R2: %f' % r2)

    return r2

# 差分
def difference(dataset):
    diff = list()
    for i in range(1, len(dataset)):
        value = dataset[i] - dataset[i - 1]
        diff.append(value)
    return diff

# 加载数据
dataset = read_csv('E:\lyf_ML_Drought\coding\ML_Drought_Prediction\indices_caculate\\result\multi_spei_csv\SPEI-12\Multi_SPEI-12_56080.txt',
    header=0, parse_dates=[['year', 'month']], index_col=0, date_parser=parse)
# 把时间当索引
dataset.index.name = 'time'
# 删掉一些列
dataset.drop(columns=['average_air_pressure',
                      'average_water_air_pressure',
                      'low_temp', 'high_temp',
                      'precipitation',
                      'temperature',
                      'humidity'], inplace=True)
X = dataset.values
X = X.astype(float)

# 对数据进行差分
stationary = difference(X) # 一阶差分
stationary = difference(stationary) # 二阶差分
# stationary.index = dataset.index[1:]
# 检查稳定性(ADF)
result = adfuller(stationary)
print('ADF Statistic: %f' % result[0])
print('p-value: %f' % result[1])
print('Critical Values:')
for key, value in result[4].items():
    print('\t%s: %.3f' % (key, value))
print('\n')
# 随机性检验（白噪声检验）
p_value = acorr_ljungbox(X, lags=1)
print('白噪音检验：', p_value[1][0])
# 若统计量的P值小于显著性水平0.05，则可以以95%的置信水平拒绝原假设，认为序列为非白噪声序列（否则，接受原假设，认为序列为纯随机序列。）
# 由于P值远小于0.05所以拒绝原假设，认为时间序列是非白噪声的，具有时间上的相关性。

# 画出差分数据
S = pandas.DataFrame(data=stationary, index=dataset.index[2:],)
S.plot()
pyplot.show()

#画自相关图和偏自相关图，定阶
pyplot.figure()
pyplot.subplot(211)
plot_acf(dataset, ax=pyplot.gca())
pyplot.subplot(212)
plot_pacf(dataset, ax=pyplot.gca())
pyplot.show()

#数据准备
train_size = int(len(X) * 0.95)
print('训练数据个数：', train_size)
print('测试数据个数：', int(len(X))-train_size)
train, test = X[0:train_size], X[train_size:]

# 前向验证
history = [x for x in train]
predictions = list()
for i in range(len(test)):
    # 预测
    model = ARIMA(history, order=(9,2,3))
    model_fit = model.fit(disp=0)
    yhat = model_fit.forecast()[0]
    predictions.append(yhat)
    # 观察
    obs = test[i]
    history.append(obs)
    print('>Predicted=%.5f, Expected=%.5f' % (yhat, obs))
# 性能
print("测试模型的评估：")
evaluate_forecasts(test, predictions)
# 画图
pyplot.plot(test)
pyplot.plot(predictions, color='red')
pyplot.show()